import numpy as np
from tqdm import tqdm

def traces_2O_pre(traces,midvalue=None,out_point_n=1000,pre_func='abs_diff'):
    if pre_func=='abs_diff':
        func=lambda a,b:abs(a-b)
    elif pre_func=='multiply':
        func = lambda a, b: a*b
    elif pre_func=='sum_square':
        func = lambda a, b:(a+b)**2
    else:
        func=lambda a,b:abs(a-b)**2
    cols_n=int((traces.shape[1]*(traces.shape[1]-1))/2)
    if midvalue is not None:
        if np.array(midvalue).ndim==2:
            midvalue=midvalue@2**np.array(range(8))[::-1]
        test_n=min(1000,len(midvalue))
        test_traces=traces[:test_n]
        test_midvalue=midvalue[:test_n].flatten()
        out_n=min(cols_n,out_point_n)
        mc_start,mc_end,mc_cof=[0]*out_n,[0]*out_n,np.zeros((out_n))
        for i in tqdm(range(test_traces.shape[1]-1)):
            for j in range(i+1,test_traces.shape[1]):
                pro_col=func(test_traces[:,i],test_traces[:,j])
                corr = abs(np.corrcoef(test_midvalue, pro_col)[0][1])
                if min(mc_cof)<corr:
                    index=np.where(mc_cof==min(mc_cof))[0][0]
                    mc_start[index],mc_end[index],mc_cof[index]=i,j,corr
        traces_pre=np.zeros((traces.shape[0],out_n))
        for i in tqdm(range(out_n)):
            traces_pre[:,i]=func(traces[:,mc_start[i]],traces[:,mc_end[i]])
    else:
        traces_pre=np.zeros((traces.shape[0],cols_n))
        n=0
        for i in tqdm(range(traces.shape[1]-1)):
            for j in range(i+1,traces.shape[1]):
                traces_pre[:,n]=func(traces[:,i],traces[:,j])
                n+=1
    return traces_pre
